import numpy as np
import torch
from matplotlib import pyplot as plt
from torch import nn
import pandas as pd
from sklearn import preprocessing
import torch.utils.data
import torch.utils.data as Data
from sklearn.metrics import mean_squared_error  # MSE
from sklearn.metrics import mean_absolute_error  # MAE

# 0. 模型的定义
# RNN类定义
class RNN(nn.Module):
    def __init__(self, series_dim, hidden_size, num_layers):
        super(RNN, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.rnn = nn.RNN(series_dim, hidden_size, num_layers, batch_first=True)
        self.linear = nn.Linear(hidden_size, series_dim)

    def __call__(self, input, state):
        hiddens, outputs = self.rnn(input, state)
        return self.linear(hiddens)

    def init_state(self, batch_size):
        return torch.zeros((self.num_layers, batch_size, self.hidden_size))


# LSTM类定义
class LSTM(nn.Module):
    def __init__(self, series_dim, hidden_size, num_layers):
        super(LSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(series_dim, hidden_size, num_layers, batch_first=True)
        self.linear = nn.Linear(hidden_size, series_dim)

    def __call__(self, input, state):
        hiddens, outputs = self.lstm(input, state)
        return self.linear(hiddens)

    def init_state(self, batch_size):
        h = torch.zeros(self.num_layers, batch_size,
                        self.hidden_size, requires_grad=True)
        c = torch.zeros(self.num_layers, batch_size,
                        self.hidden_size, requires_grad=True)
        return (h, c)


# 准备进行数据的处理（划分样本和标签，训练集与测试集的划分）
# 1. 读取csv表格数据
data = pd.read_csv('data/pollution.csv').values[6000:10000]

# 2. 选取 pollution,dew,temp,press作为我们要学习和预测的数据
data = data[:, [1, 2, 3, 4]]
plt.plot(data)
plt.show()

# 3. 数据归一化：为了消除不同维度指标之间的量纲影响，需要进行数据标准化处理，以解决数据指标之间的可比性。按均值和方差将数据限制在[-1,1]或者[0,1]
min_max_scaler = preprocessing.MinMaxScaler()
data = min_max_scaler.fit_transform(data)

# 3. 分开样本和标签 样本:(x_0,x_1,...x_n-1) -> 标签:(x_1,x_2,...x_n), 样本:(x_1,x_2,...x_n) -> 标签:(x_2,x_3,...x_n+1) ,...
WINDOW_SIZE = 50
samples = []
labels = []
for i in range(len(data) - WINDOW_SIZE - 1):
    samples.append(data[i:i + WINDOW_SIZE])
    labels.append(data[i + 1:i + WINDOW_SIZE + 1])
samples = torch.tensor(np.array(samples), dtype=torch.float32)
labels = torch.tensor(np.array(labels), dtype=torch.float32)

# 4. 划分训练集，测试集 按4:1划分
train_test_boundary = int(len(data) * 0.8)
train_samples = samples[:train_test_boundary, :]
train_labels = labels[:train_test_boundary, :]
test_samples = samples[train_test_boundary:, :]

# 5. 使用pytorch的dataloader构建数据集
BATCH_SIZE = 32
train_dataset = Data.TensorDataset(train_samples, train_labels)
train_loader = Data.DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True)

# 数据已准备好，开始训练
# 6. 实例化模型，这里使用LSTM作为我们的模型
TIME_SERIES_DIM = 4
HIDDEN_SIZE = 60
NUM_LAYERS = 1
model = LSTM(TIME_SERIES_DIM, HIDDEN_SIZE, NUM_LAYERS)

# 7. 定义损失函数和优化器等
LEARNING_RATE = 0.001
loss_function = nn.MSELoss()
updater = torch.optim.Adam(model.parameters(), LEARNING_RATE)

# 8. 开始训练
EPOCH = 10
for epoch in range(EPOCH):
    current_epoch_loss = 0
    batches_per_epoch = 0
    for X, Y in train_loader:
        # 循环神经网络初始化状态h
        state = model.init_state(BATCH_SIZE)
        # 把当前状态state和当前输入X输入进模型
        Y_pred = model(X, state)
        # 通过损失函数和优化器来更新
        updater.zero_grad()
        loss = loss_function(Y_pred, Y)
        current_epoch_loss += float(loss) * BATCH_SIZE
        batches_per_epoch += BATCH_SIZE
        loss.backward()
        updater.step()  # 梯度下降更新权重
    print(f"*Current epoch:{epoch} training loss:{current_epoch_loss / batches_per_epoch}")

# 9. 预测
predicted_data = []
for test_sample in test_samples:
    state = model.init_state(batch_size=1)
    predicted_label = model(test_sample.reshape(1, test_sample.shape[0], -1), state)
    predicted_data.append(predicted_label[0][-1].detach().numpy())
predicted_data = np.array(predicted_data)

# 10. 计算误差并绘制预测结果
true_data = data[train_test_boundary + WINDOW_SIZE + 1:]
MSE = mean_squared_error(true_data, predicted_data)
MAE = mean_absolute_error(true_data, predicted_data)
print(f"Mean squared error:{round(MSE, 5)} Mean absolute error:{round(MAE, 5)}")
plt.plot(true_data[:, 0])
plt.plot(predicted_data[:, 0])
plt.show()
